import numpy as np
import Forward_Kinematics_er7pro as fk

def srsInverseKinematics(psi, T, qPrev):
    '''
    %% @brief      由臂角和末端位姿可以得到8组机械臂关节角度,在连续轨迹规划过程中，根据关节角度变化最小原则选出一组最优关节角度
    %% @author     wangkai
    %% @version    1.0
    %% @param[in]  psi:臂角
    %% @param[in]  T:4x4齐次变换矩阵,即机械臂末端位姿矩阵
    %% @param[in]  qPrev:1x7关节角度向量,对应上一个轨迹点
    %% @return     1x7最优关节逆解
    '''
    dBS = 0.404
    dSE = 0.4375
    dEW = 0.4125
    dWT = 0.2755
    
    # 机械臂参数DH参数，详细参数参见文档
    nx, ox, ax, px = T[0]
    ny, oy, ay, py = T[1]
    nz, oz, az, pz = T[2]
    
    # 0->7旋转矩阵
   
    R07 = np.array([[nx, ox, ax],
                    [ny, oy, ay],
                    [nz, oz, az]])
    
    pSW = np.array([-dWT*ax+px, -dWT*ay+py, -dWT*az+pz-dBS])
    dSW = np.linalg.norm(pSW)
    
    # 判断是否存在逆解
    
    if dSW >= (dSE + dEW):
        print('该位姿无解')
        return None
    
    # 计算关节角
    theta_ESW_1 =  np.arccos((dSE**2+dSW**2-dEW**2)/(2*dSE*dSW))
    theta_ESW_2 = -np.arccos((dSE**2+dSW**2-dEW**2)/(2*dSE*dSW))

    theta_SEW_1 =  np.arccos((dSE**2+dEW**2-dSW**2)/(2*dSE*dEW))
    theta_SEW_2 = -np.arccos((dSE**2+dEW**2-dSW**2)/(2*dSE*dEW))
    
    q4_1 = np.pi-theta_SEW_1
    q4_2 = np.pi-theta_SEW_2
    
    
    uSW = pSW/dSW
    # uSW = uSW.squeeze()
    z0 = np.array([0, 0, 1])
    kk = np.cross(uSW,z0)/np.linalg.norm(np.cross(uSW,z0))
    # kk = kk.squeeze()
    I = np.eye(3)
    print(kk.shape)
    K = np.array([[0, -kk[2], kk[1]], 
                  [kk[2], 0, -kk[0]], 
                  [-kk[1], kk[0], 0]])
    
    # 罗德里格斯公式
    R_k_theta_ESW_1 = I + np.sin(theta_ESW_1)*K + (1-np.cos(theta_ESW_1))*K@K
    R_k_theta_ESW_2 = I + np.sin(theta_ESW_2)*K + (1-np.cos(theta_ESW_2))*K@K

    y3_1 = -R_k_theta_ESW_1@uSW
    x3_1 = ((pSW+y3_1*dSE+np.cos(q4_1)*y3_1*dEW)/(np.sin(q4_1)*dEW))
    z3_1 = np.cross(x3_1,y3_1)

    y3_2 = -R_k_theta_ESW_2@uSW
    x3_2 = ((pSW+y3_2*dSE+np.cos(q4_2)*y3_2*dEW)/(np.sin(q4_2)*dEW))
    z3_2 = np.cross(x3_2,y3_2)

    # R03_psi0_1 = np.array([x3_1, y3_1, z3_1])
    # R03_psi0_2 = np.array([x3_2, y3_2, z3_2])
    
    R03_psi0_1 = np.c_[x3_1, y3_1, z3_1]
    R03_psi0_2 = np.c_[x3_2, y3_2, z3_2]

    
    # uSW反对称矩阵
    uSW_cha = np.array([[0, -uSW[2], uSW[1]], 
                        [uSW[2], 0, -uSW[0]], 
                        [-uSW[1], uSW[0], 0]])
    R34_1 = np.array([[np.cos(q4_1), 0, np.sin(q4_1)], 
                      [np.sin(q4_1), 0, -np.cos(q4_1)], 
                      [0, 1, 0]])
    R34_2 = np.array([[np.cos(q4_2), 0, np.sin(q4_2)], 
                      [np.sin(q4_2), 0, -np.cos(q4_2)], 
                      [0, 1, 0]])
    
    # 前三个关节系数矩阵
    As_1 = uSW_cha@R03_psi0_1
    Bs_1 = -uSW_cha@uSW_cha@R03_psi0_1
    Cs_1 = (I+uSW_cha@uSW_cha)@R03_psi0_1

    As_2 = uSW_cha@R03_psi0_2
    Bs_2 = -uSW_cha@uSW_cha@R03_psi0_2
    Cs_2 = (I+uSW_cha@uSW_cha)@R03_psi0_2
    
    Aw_1 = (R34_1.T)@(As_1.T)@R07
    Bw_1 = (R34_1.T)@(Bs_1.T)@R07
    Cw_1 = (R34_1.T)@(Cs_1.T)@R07

    Aw_2 = (R34_2.T)@(As_2.T)@R07
    Bw_2 = (R34_2.T)@(Bs_2.T)@R07
    Cw_2 = (R34_2.T)@(Cs_2.T)@R07
    Cw_2 = (R34_2.T)@(Cs_2.T)@R07
    
    # arccos函数返回值为0°到180°
    q2_1_1 =  np.arccos(-As_1[2,1]*np.sin(psi)-Bs_1[2,1]*np.cos(psi)-Cs_1[2,1])
    q2_1_2 = -np.arccos(-As_1[2,1]*np.sin(psi)-Bs_1[2,1]*np.cos(psi)-Cs_1[2,1])
    q2_2_1 =  np.arccos(-As_2[2,1]*np.sin(psi)-Bs_2[2,1]*np.cos(psi)-Cs_2[2,1])
    q2_2_2 = -np.arccos(-As_2[2,1]*np.sin(psi)-Bs_2[2,1]*np.cos(psi)-Cs_2[2,1])
    
    # 关节2小于0
    q1_1_1 = np.atan2(As_1[1,1]*np.sin(psi)+Bs_1[1,1]*np.cos(psi)+Cs_1[1,1],As_1[0,1]*np.sin(psi)+Bs_1[0,1]*np.cos(psi)+Cs_1[0,1])
    q1_2_1 = np.atan2(As_2[1,1]*np.sin(psi)+Bs_2[1,1]*np.cos(psi)+Cs_2[1,1],As_2[0,1]*np.sin(psi)+Bs_2[0,1]*np.cos(psi)+Cs_2[0,1])

    q3_1_1 = np.atan2(-As_1[2,2]*np.sin(psi)-Bs_1[2,2]*np.cos(psi)-Cs_1[2,2],As_1[2,0]*np.sin(psi)+Bs_1[2,0]*np.cos(psi)+Cs_1[2,0])
    q3_2_1 = np.atan2(-As_2[2,2]*np.sin(psi)-Bs_2[2,2]*np.cos(psi)-Cs_2[2,2],As_2[2,0]*np.sin(psi)+Bs_2[2,0]*np.cos(psi)+Cs_2[2,0])

    # 关节2大于0
    q1_1_2 = np.atan2(-As_1[1,1]*np.sin(psi)-Bs_1[1,1]*np.cos(psi)-Cs_1[1,1],-As_1[0,1]*np.sin(psi)-Bs_1[0,1]*np.cos(psi)-Cs_1[0,1])
    q1_2_2 = np.atan2(-As_2[1,1]*np.sin(psi)-Bs_2[1,1]*np.cos(psi)-Cs_2[1,1],-As_2[0,1]*np.sin(psi)-Bs_2[0,1]*np.cos(psi)-Cs_2[0,1])

    q3_1_2 = np.atan2(As_1[2,2]*np.sin(psi)+Bs_1[2,2]*np.cos(psi)+Cs_1[2,2],-As_1[2,0]*np.sin(psi)-Bs_1[2,0]*np.cos(psi)-Cs_1[2,0])
    q3_2_2 = np.atan2(As_2[2,2]*np.sin(psi)+Bs_2[2,2]*np.cos(psi)+Cs_2[2,2],-As_2[2,0]*np.sin(psi)-Bs_2[2,0]*np.cos(psi)-Cs_2[2,0])

    q6_1_1 =  np.arccos(Aw_1[2,2]*np.sin(psi)+Bw_1[2,2]*np.cos(psi)+Cw_1[2,2])
    q6_1_2 = -np.arccos(Aw_1[2,2]*np.sin(psi)+Bw_1[2,2]*np.cos(psi)+Cw_1[2,2])
    q6_2_1 =  np.arccos(Aw_2[2,2]*np.sin(psi)+Bw_2[2,2]*np.cos(psi)+Cw_2[2,2])
    q6_2_2 = -np.arccos(Aw_2[2,2]*np.sin(psi)+Bw_2[2,2]*np.cos(psi)+Cw_2[2,2])

    # 关节6小于0
    q5_1_1 = np.atan2(-Aw_1[1,2]*np.sin(psi)-Bw_1[1,2]*np.cos(psi)-Cw_1[1,2],-Aw_1[0,2]*np.sin(psi)-Bw_1[0,2]*np.cos(psi)-Cw_1[0,2])
    q5_2_1 = np.atan2(-Aw_2[1,2]*np.sin(psi)-Bw_2[1,2]*np.cos(psi)-Cw_2[1,2],-Aw_2[0,2]*np.sin(psi)-Bw_2[0,2]*np.cos(psi)-Cw_2[0,2])

    q7_1_1 = np.atan2(-Aw_1[2,1]*np.sin(psi)-Bw_1[2,1]*np.cos(psi)-Cw_1[2,1],Aw_1[2,0]*np.sin(psi)+Bw_1[2,0]*np.cos(psi)+Cw_1[2,0])
    q7_2_1 = np.atan2(-Aw_2[2,1]*np.sin(psi)-Bw_2[2,1]*np.cos(psi)-Cw_2[2,1],Aw_2[2,0]*np.sin(psi)+Bw_2[2,0]*np.cos(psi)+Cw_2[2,0])

    # 关节6大于0
    q5_1_2 = np.atan2(Aw_1[1,2]*np.sin(psi)+Bw_1[1,2]*np.cos(psi)+Cw_1[1,2],Aw_1[0,2]*np.sin(psi)+Bw_1[0,2]*np.cos(psi)+Cw_1[0,2])
    q5_2_2 = np.atan2(Aw_2[1,2]*np.sin(psi)+Bw_2[1,2]*np.cos(psi)+Cw_2[1,2],Aw_2[0,2]*np.sin(psi)+Bw_2[0,2]*np.cos(psi)+Cw_2[0,2])

    q7_1_2 = np.atan2(Aw_1[2,1]*np.sin(psi)+Bw_1[2,1]*np.cos(psi)+Cw_1[2,1],-Aw_1[2,0]*np.sin(psi)-Bw_1[2,0]*np.cos(psi)-Cw_1[2,0])
    q7_2_2 = np.atan2(Aw_2[2,1]*np.sin(psi)+Bw_2[2,1]*np.cos(psi)+Cw_2[2,1],-Aw_2[2,0]*np.sin(psi)-Bw_2[2,0]*np.cos(psi)-Cw_2[2,0])

    # 给出8组解
    q_all = np.array([[q1_1_1, q2_1_2, q3_1_1, q4_1, q5_1_1, q6_1_2, q7_1_1],
                      [q1_1_1, q2_1_2, q3_1_1, q4_1, q5_1_2, q6_1_1, q7_1_2],
                      [q1_1_2, q2_1_1, q3_1_2, q4_1, q5_1_1, q6_1_2, q7_1_1],
                      [q1_1_2, q2_1_1, q3_1_2, q4_1, q5_1_2, q6_1_1, q7_1_2],
                      [q1_2_1, q2_2_2, q3_2_1, q4_2, q5_2_1, q6_2_2, q7_2_1],
                      [q1_2_1, q2_2_2, q3_2_1, q4_2, q5_2_2, q6_2_1, q7_2_2],
                      [q1_2_2, q2_2_1, q3_2_2, q4_2, q5_2_1, q6_2_2, q7_2_1],
                      [q1_2_2, q2_2_1, q3_2_2, q4_2, q5_2_2, q6_2_1, q7_2_2]])
    
    valid_q_all = []
    validNum = 0
    for i in range(8):
        curQ = q_all[i, :]
        if (not np.isnan(curQ).any()) and is_in_joint_limit(curQ):
            valid_q_all.append(curQ)
            validNum += 1

    if validNum == 0:
        print("该位姿和臂角下机械臂无关节逆解")
        return None

    print(f'该位姿和臂角下机械臂有效逆解个数: {validNum}')
    res = valid_q_all[0]
    variationAngle = np.sum(np.abs(qPrev[:6] - valid_q_all[0][:6]))

    for i in range(validNum):
        curVariationAngleCur = np.sum(np.abs(qPrev[:6] - valid_q_all[i][:6]))
        if curVariationAngleCur <= variationAngle:
            variationAngle = curVariationAngleCur
            res = valid_q_all[i]

    if len(res) > 0:
        if any(np.abs(qPrev[:7] - res[:7]) * 180 / np.pi > 175):
            print('逆解存在突变，该位姿和臂角下机械臂无关节逆解')
            return None
    return res


def is_in_joint_limit(q):
    """
    判断机械臂关节角度是否位于关节极限
    """
    if (q[0] > np.deg2rad(170)) or (q[0] < -np.deg2rad(170)):
        return False

    if (q[1] > np.deg2rad(120)) or (q[1] < -np.deg2rad(120)):
        return False

    if (q[2] > np.deg2rad(170)) or (q[2] < -np.deg2rad(170)):
        return False

    if (q[3] > np.deg2rad(120)) or (q[3] < -np.deg2rad(120)):
        return False

    if (q[4] > np.deg2rad(170)) or (q[4] < -np.deg2rad(170)):
        return False

    if (q[5] > np.deg2rad(120)) or (q[5] < -np.deg2rad(120)):
        return False

    if (q[6] > np.deg2rad(360)) or (q[6] < -np.deg2rad(360)):
        return False
    return True

if __name__ == '__main__':
    psi = -np.pi/2
    T = np.array([[0, 1, 0, 0.5],
                  [0, 1, 0, 0],
                  [1, 0, 0, -0.5],
                  [0, 0, 0, 1]])
    qPrev = np.array([0, 0, 0, 0, 0, 0, 0])
    q = srsInverseKinematics(psi, T, qPrev)
    if q is not None:
        bb = fk.srs_forward_kinematics(q)
        print(f"1:{bb}")